import os
import dgl
import json
import torch
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from sklearn.metrics import f1_score

from configs.Params import Params
from nets.load_net import gnn_model
from data_process import data_process
from farthest_insertion import run_insertion


def gpu_setup(use_gpu, gpu_id):
    """GPU设置"""
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    if torch.cuda.is_available() and use_gpu:
        print('cuda available with GPU:', torch.cuda.get_device_name(0))
        device = torch.device("cuda")
    else:
        print('cuda not available')
        device = torch.device("cpu")
    return device


def read_data(data_path):
    graph_list = []
    edge_labels = []
    # 读取json文件
    with open(data_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    for i in range(len(data)):
        # 数据预处理
        i = str(i)
        num_nodes = len(data[i]['ndata'])
        edges = torch.tensor(data[i]['edges'], dtype=torch.int64)
        n_feat = torch.tensor(data[i]['ndata'], dtype=torch.float32)
        e_feat = torch.tensor(data[i]['edata'], dtype=torch.float32).view(-1, 1)
        edge_label = data[i]['label']
        # 构造图
        graph = dgl.DGLGraph()
        graph.add_nodes(num_nodes)
        graph.add_edges(edges[:, 0], edges[:, 1])
        graph.ndata['feat'] = n_feat
        graph.edata['feat'] = e_feat
        # 保存数据
        graph_list.append(graph)
        edge_labels.append(edge_label)

    return graph_list, edge_labels

def plat_TSP(graph, node_pos, edges_list, tour_list, title="tour", show=True):
    """给定边集合画图"""
    plt.figure()
    plt.title(title)
    graph = graph.to_networkx()
    graph = graph.to_undirected()
    pos = dict(zip(range(len(node_pos)), node_pos))
    # 控制点的大小
    if len(node_pos) >= 200:
        node_size = 50 / (len(node_pos) // 50)
    else:
        node_size = 100
    nx.draw_networkx_nodes(graph, pos, node_size=node_size, label=True)
    # 画边
    # alpha：透明度参数；width：宽度参数
    nx.draw_networkx_edges(graph, pos, edgelist=edges_list, alpha=0.25, width=0.1)
    nx.draw_networkx_edges(graph, pos, edgelist=tour_list, alpha=1, width=1.5, edge_color='r')
    if show:
        plt.show()


def plat_tour(graph, labels, title="tour", show=True):
    """给定边标签画图"""
    # 获得每个node的位置信息
    node_pos = graph.ndata['feat']
    # 转为list
    node_pos = node_pos.numpy().tolist()
    # 将所有的边连线转为tuple类型
    edges = graph.edges()
    u = edges[0].numpy()
    v = edges[1].numpy()
    edges_list = []
    for i in range(len(u)):
        edges_list.append((u[i], v[i]))
    tour_list = []
    for i in range(len(labels)):
        if labels[i] == 1:
            tour_list.append(edges_list[i])
    plat_TSP(graph, node_pos, edges_list, tour_list, title, show)


def plat_compare(graph, pred_labels, real_labels):
    """画预测路由和真实路由图以进行比较"""
    plat_tour(graph, pred_labels, title="pred_tour", show=False)
    plat_tour(graph, real_labels, title="truth_tour", show=False)
    plt.show()


def binary_f1_score(scores, targets):
    """计算F1值"""
    y_true = targets.cpu().numpy()
    y_pred = scores.argmax(dim=1).cpu().numpy()
    return f1_score(y_true, y_pred, average='binary')


def evaluate_net(model, device, graph, label):
    """验证/测试已训练好的模型"""
    model.eval()
    with torch.no_grad():
        graph = graph.to(device)
        x = graph.ndata['feat'].to(device)
        e = graph.edata['feat'].to(device)
        label = label.to(device)
        score = model.forward(graph, x, e)
        f1_value = binary_f1_score(score, label)
        # 求预测边与真实边
        score = torch.softmax(score, dim=-1)
        score = score.cpu().numpy()
        pred_labels = np.argmax(score, axis=1)
        real_labels = label.cpu().numpy()
    return pred_labels, real_labels, f1_value


def read_checkpoint(model_name, graph, label, checkpoint_path):
    """读取checkpoint以测试模型"""
    device = gpu_setup(True, 0)
    params = Params(model_name, device)
    # 定义模型
    model = gnn_model(model_name, params)
    model = model.to(device)
    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
    pred_labels, real_labels, f1_value = evaluate_net(model, device, graph, label)
    print("F1_Value: ", f1_value)
    plat_compare(graph, pred_labels, real_labels)



def idx_data_test(model_name, data_path, checkpoint_path, idx):
    """测试 test数据集中的 idx号数据"""
    # 读取数据
    graph_list, edge_labels = read_data(data_path)
    graph = graph_list[idx]
    label = edge_labels[idx]
    label = torch.LongTensor(label)  # 转为LongTensor
    # 读取checkpoint来测试
    read_checkpoint(model_name, graph, label, checkpoint_path)


def random_test(model_name, points, checkpoint_path):
    """测试一个随机数据"""
    # 使用最远插入求近似解
    cost, tour = run_insertion(points)
    # 数据预处理
    graph, label = data_process(points, tour, self_loop=False, num_near=None)
    # 读取checkpoint来测试
    read_checkpoint(model_name, graph, label, checkpoint_path)


if __name__ == '__main__':
    # 测试模型
    # model_name = 'GCN'
    # model_name = 'GAT'
    model_name = 'GatedGCN'
    checkpoint_path = "../checkpoints/" + str(model_name) + "/100k.pkl"
    # 读取数据
    data_path = "../dataset/small/test_data.json"
    idx_data_test(model_name, data_path, checkpoint_path, idx=23)

    num_city = 30
    points = np.random.uniform(0, 1, size=(num_city, 2))
    random_test(model_name, points, checkpoint_path)
